import pymysql
from pymysql import cursors

from utils import Config, logger


class MySQL:
    """MySQL工具类"""

    def __init__(
            self,
            username=None,  # The first four arguments is based on DB-API 2.0 recommendation.
            password="",
            host=None,
            database=None,
            port=0,
            charset=""
    ):
        self.connect = pymysql.connect(host=host, port=int(port),
                                       user=username, password=password,
                                       database=database, charset=charset)
        self.cursor = self.connect.cursor(cursors.DictCursor)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.cursor.close()
        self.connect.close()

    def query(self, query: str, args=None):
        try:
            self.cursor.execute(query, args)
            return self.cursor.fetchall()
        except pymysql.MySQLError as e:
            logger.error(e)
            self.connect.rollback()


if __name__ == '__main__':
    mysql_conf = Config().get_mysql_config()
    with MySQL(**mysql_conf) as mysql:
        print(mysql.query('select * from info'))
